Java ASR Module binding #16979
Java ASR Module binding #16979
Conversation
meta-cla
bot
added
the
CLA Signed
There was a problem hiding this comment.
Pull request overview
This pull request adds Java/Kotlin bindings for the ExecuTorch ASR (Automatic Speech Recognition) module, enabling Android applications to use ASR models like Whisper. The implementation follows a similar pattern to the existing LLM module bindings.
Changes:
- Added JNI layer implementation (
jni_layer_asr.cpp) to bridge C++ ASR runner with Java/Kotlin - Created Kotlin API classes:
AsrModule,AsrCallback, andAsrTranscribeConfigfor Android integration - Updated build scripts and CMake configuration to include ASR runner support
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 15 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/build_android_library.sh | Adds EXECUTORCH_BUILD_EXTENSION_ASR_RUNNER flag to Android build configuration |
| extension/android/jni/jni_layer_asr.cpp | Implements JNI bindings for ASR runner including native methods and callback handling |
| extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrTranscribeConfig.kt | Defines configuration data class for ASR transcription parameters |
| extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt | Main Kotlin API class for ASR module with transcription methods |
| extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrCallback.kt | Callback interface for receiving transcription tokens and completion events |
| extension/android/CMakeLists.txt | Updates CMake to include ASR JNI layer in the build |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Outdated
Show resolved
Hide resolved
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Show resolved
Hide resolved
...android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrCallback.kt
Show resolved
Hide resolved
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Show resolved
Hide resolved
Use int64_t arithmetic to detect overflow when computing batchSize * timeSteps * featureDim before casting to jsize. This prevents silent overflow that could cause incorrect validation and potential out-of-bounds memory access.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if (scopedCallback) { | ||
| initCallbackCache(env); | ||
|
|
||
| jobject callbackRef = scopedCallback.get(); | ||
| tokenCallback = [env, callbackRef, &tokenBuffer](const std::string& token) { | ||
| tokenBuffer += token; | ||
| if (!utf8_check_validity(tokenBuffer.c_str(), tokenBuffer.size())) { | ||
| ET_LOG( | ||
| Info, "Current token buffer is not valid UTF-8. Waiting for more."); | ||
| return; | ||
| } | ||
|
|
||
| std::string completeToken = tokenBuffer; | ||
| tokenBuffer.clear(); | ||
|
|
||
| jstring jToken = env->NewStringUTF(completeToken.c_str()); | ||
| env->CallVoidMethod(callbackRef, callbackCache.onTokenMethod, jToken); | ||
| if (env->ExceptionCheck()) { | ||
| ET_LOG(Error, "Exception occurred in AsrCallback.onToken"); | ||
| env->ExceptionClear(); | ||
| } | ||
| env->DeleteLocalRef(jToken); | ||
| }; | ||
| } |
There was a problem hiding this comment.
After calling initCallbackCache, there's no validation that the callback cache was properly initialized. If initCallbackCache failed to find the class or methods, callbackCache.onTokenMethod and callbackCache.onCompleteMethod will be nullptr, causing crashes at lines 326 and 342. Add validation checks after initCallbackCache to ensure the cache is properly initialized before using the callback.
| // Helper to get a string from jstring | ||
| std::string jstringToString(JNIEnv* env, jstring jstr) { | ||
| if (jstr == nullptr) { | ||
| return ""; | ||
| } | ||
| const char* chars = env->GetStringUTFChars(jstr, nullptr); | ||
| std::string result(chars); | ||
| env->ReleaseStringUTFChars(jstr, chars); | ||
| return result; | ||
| } |
There was a problem hiding this comment.
The jstringToString helper function is a simple utility that could potentially be shared across JNI layers. Consider adding it to jni/jni_helper.h if it's used in multiple places. However, this is not critical given its simplicity.
| jobject callbackRef = scopedCallback.get(); | ||
| tokenCallback = [env, callbackRef, &tokenBuffer](const std::string& token) { | ||
| tokenBuffer += token; | ||
| if (!utf8_check_validity(tokenBuffer.c_str(), tokenBuffer.size())) { | ||
| ET_LOG( | ||
| Info, "Current token buffer is not valid UTF-8. Waiting for more."); | ||
| return; | ||
| } | ||
|
|
||
| std::string completeToken = tokenBuffer; | ||
| tokenBuffer.clear(); | ||
|
|
||
| jstring jToken = env->NewStringUTF(completeToken.c_str()); | ||
| env->CallVoidMethod(callbackRef, callbackCache.onTokenMethod, jToken); | ||
| if (env->ExceptionCheck()) { | ||
| ET_LOG(Error, "Exception occurred in AsrCallback.onToken"); | ||
| env->ExceptionClear(); | ||
| } | ||
| env->DeleteLocalRef(jToken); | ||
| }; |
There was a problem hiding this comment.
The lambda captures JNIEnv* by value, which is only safe if the callback is guaranteed to be invoked synchronously on the same thread. JNIEnv* pointers are thread-local in JNI and cannot be safely used across threads. While the current runner implementation appears to invoke callbacks synchronously, this creates a fragile dependency. Consider documenting this requirement explicitly or adding a runtime check to verify thread safety.
| jobject callbackRef = scopedCallback.get(); | |
| tokenCallback = [env, callbackRef, &tokenBuffer](const std::string& token) { | |
| tokenBuffer += token; | |
| if (!utf8_check_validity(tokenBuffer.c_str(), tokenBuffer.size())) { | |
| ET_LOG( | |
| Info, "Current token buffer is not valid UTF-8. Waiting for more."); | |
| return; | |
| } | |
| std::string completeToken = tokenBuffer; | |
| tokenBuffer.clear(); | |
| jstring jToken = env->NewStringUTF(completeToken.c_str()); | |
| env->CallVoidMethod(callbackRef, callbackCache.onTokenMethod, jToken); | |
| if (env->ExceptionCheck()) { | |
| ET_LOG(Error, "Exception occurred in AsrCallback.onToken"); | |
| env->ExceptionClear(); | |
| } | |
| env->DeleteLocalRef(jToken); | |
| }; | |
| JavaVM* jvm = nullptr; | |
| if (env->GetJavaVM(&jvm) != JNI_OK || jvm == nullptr) { | |
| ET_LOG(Error, "Failed to get JavaVM; token callback will be disabled."); | |
| } else { | |
| jobject callbackRef = scopedCallback.get(); | |
| tokenCallback = | |
| [jvm, callbackRef, &tokenBuffer](const std::string& token) { | |
| JNIEnv* envLocal = nullptr; | |
| bool detach = false; | |
| jint getEnvResult = | |
| jvm->GetEnv(reinterpret_cast<void**>(&envLocal), JNI_VERSION_1_6); | |
| if (getEnvResult == JNI_EDETACHED) { | |
| if (jvm->AttachCurrentThread(&envLocal, nullptr) != JNI_OK) { | |
| ET_LOG( | |
| Error, | |
| "Failed to attach current thread to JVM in token callback."); | |
| return; | |
| } | |
| detach = true; | |
| } else if (getEnvResult != JNI_OK || envLocal == nullptr) { | |
| ET_LOG( | |
| Error, | |
| "Failed to get JNIEnv for current thread in token callback."); | |
| return; | |
| } | |
| tokenBuffer += token; | |
| if (!utf8_check_validity(tokenBuffer.c_str(), tokenBuffer.size())) { | |
| ET_LOG( | |
| Info, | |
| "Current token buffer is not valid UTF-8. Waiting for more."); | |
| if (detach) { | |
| jvm->DetachCurrentThread(); | |
| } | |
| return; | |
| } | |
| std::string completeToken = tokenBuffer; | |
| tokenBuffer.clear(); | |
| jstring jToken = envLocal->NewStringUTF(completeToken.c_str()); | |
| envLocal->CallVoidMethod( | |
| callbackRef, callbackCache.onTokenMethod, jToken); | |
| if (envLocal->ExceptionCheck()) { | |
| ET_LOG(Error, "Exception occurred in AsrCallback.onToken"); | |
| envLocal->ExceptionClear(); | |
| } | |
| envLocal->DeleteLocalRef(jToken); | |
| if (detach) { | |
| jvm->DetachCurrentThread(); | |
| } | |
| }; | |
| } |
This reverts commit cf116b7.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| * | ||
| * @param token The decoded text token | ||
| */ | ||
| fun onToken(token: String) |
There was a problem hiding this comment.
Missing @DoNotStrip annotation on the onToken method. This annotation is required for JNI callback methods to prevent ProGuard/R8 from stripping them during release builds. The LlmCallback interface uses this annotation (see extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmCallback.java:28). You'll need to add the import: import com.facebook.jni.annotations.DoNotStrip
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Show resolved
Hide resolved
...android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrCallback.kt
Show resolved
Hide resolved
...n/android/executorch_android/src/main/java/org/pytorch/executorch/extension/asr/AsrModule.kt
Outdated
Show resolved
Hide resolved
Co-authored-by: Copilot <[email protected]>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| auto deleter = [env](jobject ref) { | ||
| if (ref != nullptr) { | ||
| env->DeleteGlobalRef(ref); | ||
| } | ||
| }; | ||
| jobject globalRef = obj ? env->NewGlobalRef(obj) : nullptr; |
There was a problem hiding this comment.
The deleter function captures the JNIEnv pointer by value, which is thread-local and should not be used from the deleter. JNIEnv pointers are thread-specific and cannot be safely used from a different thread context or after the current thread detaches from the JVM.
The deleter should use the JavaVM to get a thread-local JNIEnv when it needs to delete the global reference. Alternatively, since this function is only used within nativeTranscribe and the global reference is deleted before the function returns (while still on the same thread), the deleter could simply be refactored to not capture env at all and instead get it from the JavaVM, or the global reference could be deleted manually before returning instead of using a unique_ptr with a custom deleter.
| auto deleter = [env](jobject ref) { | |
| if (ref != nullptr) { | |
| env->DeleteGlobalRef(ref); | |
| } | |
| }; | |
| jobject globalRef = obj ? env->NewGlobalRef(obj) : nullptr; | |
| JavaVM* vm = nullptr; | |
| if (env != nullptr) { | |
| jint getVmResult = env->GetJavaVM(&vm); | |
| if (getVmResult != JNI_OK) { | |
| vm = nullptr; | |
| } | |
| } | |
| auto deleter = [vm](jobject ref) { | |
| if (ref == nullptr || vm == nullptr) { | |
| return; | |
| } | |
| JNIEnv* env_local = nullptr; | |
| jint envResult = vm->GetEnv(reinterpret_cast<void**>(&env_local), JNI_VERSION_1_6); | |
| bool didAttach = false; | |
| if (envResult == JNI_EDETACHED) { | |
| if (vm->AttachCurrentThread(&env_local, nullptr) != JNI_OK) { | |
| return; | |
| } | |
| didAttach = true; | |
| } else if (envResult != JNI_OK || env_local == nullptr) { | |
| return; | |
| } | |
| env_local->DeleteGlobalRef(ref); | |
| if (didAttach) { | |
| vm->DetachCurrentThread(); | |
| } | |
| }; | |
| jobject globalRef = (env != nullptr && obj != nullptr) ? env->NewGlobalRef(obj) : nullptr; |
Summary
Add java binding for extension/asr/runner/runner.h
Test plan
CI